import os, sys, copy

parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 5)))
sys.path.insert(0, parent_path)

import traceback, threading, subprocess, shutil, cv2
import numpy as np
import time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.logger import logging
from collections import deque
from queue import Queue



class VideoWriter(BaseProcessor):
    def __init__(
        self,
        keys,
        VideoWriter_save_path=None,
        VideoWriter_size=None,
        VideoWriter_fourcc="mp4v",
        VideoWriter_fps=30,
    ) -> None:
        """
        keys:
            in: 保存的单张图片
            finish_tag: 触发结束的databox中的key

        """
        self.keys = keys
        self.save_path = VideoWriter_save_path
        self.size = VideoWriter_size
        self.fourcc = VideoWriter_fourcc
        self.fps = VideoWriter_fps
        self.writer = None

    def _init_check(self):
        assert "in" in self.keys
        assert "finish_tag" in self.keys
        assert not self.size is None
        assert not self.save_path is None
        assert not self.writer is None

    def _set(self, data: Dict):
        if data.get("VideoWriter_save_path", False):
            self.save_path = data["VideoWriter_save_path"]
        if data.get("VideoWriter_fourcc", False):
            self.fourcc = data["VideoWriter_fourcc"]
        fourcc = cv2.VideoWriter_fourcc(*self.fourcc)
        # 创建对应的工作目录
        save_dir = os.path.dirname(self.save_path)
        if save_dir and not os.path.exists(save_dir):
            os.makedirs(save_dir)
        self.writer = cv2.VideoWriter(self.save_path, fourcc, self.fps, self.size)
        assert self.writer.isOpened(), "create VideoWriter fail in VideoWriter.set()"

    def _call(self, data: Dict) -> Dict:
        if data.get(self.keys["finish_tag"], False):
            self.writer.release()
        else:
            frame = data[self.keys["in"]]
            self.writer.write(cv2.resize(frame, self.size))
        return data

    def start(self):
        return super().start()

class VideoTriggerWriter(BaseProcessor):
    def __init__(
        self,
        keys,
        VideoTriggerWriter_save_dir="data/output/trigger_video",
        VideoTriggerWriter_duration=10,
        VideoTriggerWriter_size: List[int] = [1280,720],
        VideoTriggerWriter_fps=15,
    ) -> None:
        """
        keys:
            in: 保存的单张图片
            out: 保存的短视频路径
            trigger_tag: 触发录制的databox中的key
            finish_tag: 当保存结束时向databox中的保存key
        """
        self.keys = keys
        self.save_dir = VideoTriggerWriter_save_dir
        self.duration = VideoTriggerWriter_duration
        self.size = VideoTriggerWriter_size
        self.fps = int(VideoTriggerWriter_fps)
        self.img_buffer = deque(maxlen=self.fps*self.duration)
        self.save_buffer = Queue(maxsize=99)

    def _init_check(self):
        assert "in" in self.keys
        assert "trigger_tag" in self.keys
        assert not self.size is None

    def _set(self, data: Dict):
        if data.get("VideoTriggerWriter_save_dir", False):
            self.save_dir = data["VideoTriggerWriter_save_dir"]
        if data.get("VideoTriggerWriter_duration", False):
            self.duration = data["VideoTriggerWriter_duration"]
        if data.get("VideoTriggerWriter_size", False):
            self.size = data["VideoTriggerWriter_size"]
        if data.get("VideoTriggerWriter_fps", False):
            self.fps = data["VideoTriggerWriter_fps"]

    def _call(self, data: Dict) -> Dict:
        frame = data[self.keys["in"]]
        self.img_buffer.append(cv2.resize(frame,self.size))
        if data.get(self.keys["trigger_tag"], False):
            th_1 = threading.Thread(target=self.save_video)
            th_1.start()
        try:
            updata = self.save_buffer.get_nowait()
            data.update(updata)
        except:
            pass
        return data

    def save_video(self):
        time_str = time.strftime('%Y-%m-%d_%H-%M-%S',time.localtime())
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)
        save_path = self.save_dir+"/{}.mp4".format(time_str)  # 输出文件路径
        ffmpeg_cmd = [
            "ffmpeg",
            "-y",  # 自动覆盖现有文件
            "-f", "rawvideo",  # 输入格式为原始视频流
            "-vcodec", "rawvideo",  # 输入编码为原始视频流
            "-s", f"{self.size[0]}x{self.size[1]}",  # 设置输入分辨率
            "-pix_fmt", "bgr24",  # 像素格式为 RGB
            "-r", str(self.fps),  # 帧率
            "-i", "pipe:0",  # 从标准输入（stdin）读取视频
            "-vcodec", "libx264",  # 输出视频编码为 h264
            "-pix_fmt", "yuv420p",  # 输出像素格式为 yuv420p，确保兼容性
            save_path
            ]
        process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE)
        time.sleep(self.duration//2) # 等待数据载入
        save_img_list = list(self.img_buffer)
        try:
            for img_ in save_img_list:
                process.stdin.write(img_.tobytes())
        except Exception as e:
                logging.error(traceback.format_exc())
                shutil.rmtree(save_path)
        finally:
            # 关闭 stdin 并等待 ffmpeg 完成
            process.stdin.close()
            process.wait()
        
        if not os.path.exists(save_path):
            return

        try:
            self.save_buffer.put_nowait({
                self.keys["out"]:save_path,
                self.keys["finish_tag"]:True,
            })
        except:
            logging.error("VideoTriggerWriter 保存队列已满, 无法保存")
            shutil.rmtree(save_path)
                


if __name__ == "__main__":
    import yaml
    def video_writer():
        cfg_file = "/home/smartgis/workspace/project/smart_server/projects/smoke/configs/video_writer.yml"
        with open(cfg_file, "r") as f:
            cfg = yaml.safe_load(f)
        # 取对应的内部配置
        cfg = cfg["VideoWriter"]
        video_processor = VideoWriter(**cfg)
        video_processor.set(cfg)
        video_processor.start()
        # cfg[cfg["keys"]["in"]] = {}
        # video_processor(cfg)
        video_processor.close()

    def video_trigger():
        from core.task.modules.processors.steam_pull import SteamPull

        steam_puller = SteamPull(
            keys={"out":"img_data" },
            SteamPull_size=[1280,720],
            SteamPull_url= "rtmp://218.94.155.86:1936/live/livestream_74_2",
        )
        trigger_writer = VideoTriggerWriter(
            keys={
                "in":"img_data",
                "out":"trigger_video_path",
                "trigger_tag":"need_writer",
                "finish_tag":"trigger_write_done",
            },
            VideoTriggerWriter_save_dir="data/output",
            VideoTriggerWriter_fps = 15,
        )
        steam_puller.set({})
        steam_puller.start()
        trigger_writer.start()
        time.sleep(2)
        s_t = time.time()
        check = False
        while True:
            data = {}
            if not check :
                if time.time()-s_t>5:
                    data["need_writer"] = True
                    s_t = time.time()
                    check = True
            data = steam_puller(data)
            data = trigger_writer(data)
            if data.get("trigger_video_path",False):
                print(data["trigger_video_path"])
            time.sleep(1/15)

    # video_writer()
    # video_trigger()

    
    
